﻿#include "page.hpp"

extern std::map<uint64_t, Atomic<std::array<uint8_t, 64>>> global_ram;
constexpr uint8_t valid = uint16_t(PMETA::RO | PMETA::WR | PMETA::EX);

template <int N>
PageList *PageBlock<N>::get(uint64_t addr)
{
	if (idx.find(addr) == idx.end())
		return nullptr;
	auto &it = idx[addr];
	lines.splice(lines.begin(), lines, it);
	idx[addr] = lines.begin();
	return &(*lines.begin());
}
template <int N>
void PageBlock<N>::put(uint64_t addr, PageList line)
{
	Item *victim = nullptr;
	if (idx.find(addr) != idx.end()) // update
	{
		auto it = idx[addr];
		victim = &(it->second);
		lines.splice(lines.begin(), lines, it);
		idx[addr] = lines.begin();
		*it=line;
	}
	else // insert
	{
		lines.push_front(line);
		idx[addr] = lines.begin();
		if (lines.size() > N)
		{
			victim = &(lines.back().second);
			idx.erase(lines.back().first);
			lines.pop_back();
		}
	}
	if (victim != nullptr)
	{
		auto [value, paddr] = *victim;
		uint64_t target = uint64_t(value);
		global_ram[paddr >> 6].w_lock();
		memcpy(global_ram[paddr >> 6].get().data() + (paddr & 0x3f), &target, 8);
		global_ram[paddr >> 6].w_unlock();
	}
}
template <int N>
void PageBlock<N>::remove(uint64_t addr)
{
	if (idx.find(addr) == idx.end())
		return;
	auto &it = idx[addr];
	Item victim = it->second;
	lines.erase(it);
	idx.erase(addr);
	auto [value, paddr] = victim;
	uint64_t temp = uint64_t(value);
	global_ram[paddr >> 6].w_lock();
	memcpy(global_ram[paddr >> 6].get().data() + (paddr & 0x3f), &temp, 8);
	global_ram[paddr >> 6].w_unlock();
}
template <int N>
void PageBlock<N>::clear()
{
	for (auto &[addr, it] : idx)
		this->remove(addr);
}

PageList *MMU::InsSearch(uint64_t vpfn)
{ // We only find the item, you should check it in the main function stream
	uint8_t idx = vpfn & 0x0f;
	uint64_t tag = vpfn >> 4;
	if (l1i_4k.find(idx) != l1i_4k.end()) // TLB Block hit
	{
		auto it = l1i_4k[idx].get(tag);
		if (it != nullptr) // found item in l1i_4k;
		{
			tlbMeta = {false, 0};
			return it;
		}
	}
	idx = vpfn & 0x7f, tag = vpfn >> 7;
	if (l2_4k.find(idx) != l2_4k.end())
	{
		auto it = l2_4k[idx].get(tag);
		if (it != nullptr)
		{
			tlbMeta = {true, 0};
			return it;
		}
	}
	vpfn >>= 9;
	idx = vpfn & 0x0f, tag = vpfn >> 4;
	if (l1i_2m.find(idx) != l1i_2m.end())
	{
		auto it = l1i_2m[idx].get(tag);
		if (it != nullptr)
		{
			tlbMeta = {false, 1};
			return it;
		}
	}
	idx = vpfn & 0x7f, tag = vpfn >> 7;
	if (l2_2m.find(idx) != l2_2m.end())
	{
		auto it = l2_2m[idx].get(tag);
		if (it != nullptr)
		{
			tlbMeta = {true, 1};
			return it;
		}
	}
	return nullptr;
}
PageList *MMU::DataSearch(uint64_t vpfn)
{ // We only find the item, you should check it in the main function stream
	uint8_t idx = vpfn & 0x0f;
	uint64_t tag = vpfn >> 4;
	if (l1d_4k.find(idx) != l1d_4k.end())
	{
		auto it = l1d_4k[idx].get(tag);
		if (it != nullptr)
		{
			tlbMeta = {false, 0};
			return it;
		}
	}
	idx = vpfn & 0x7f, tag = vpfn >> 7;
	if (l2_4k.find(idx) != l2_4k.end())
	{
		auto it = l2_4k[idx].get(tag);
		if (it != nullptr)
		{
			tlbMeta = {true, 0};
			return it;
		}
	}
	vpfn >>= 9;
	idx = vpfn & 0x0f, tag = vpfn >> 4;
	if (l1d_2m.find(idx) != l1d_2m.end())
	{
		auto it = l1d_2m[idx].get(tag);
		if (it != nullptr)
		{
			tlbMeta = {false, 1};
			return it;
		}
	}
	idx = vpfn & 0x7f, tag = vpfn >> 7;
	if (l2_2m.find(idx) != l2_2m.end())
	{
		auto it = l2_2m[idx].get(tag);
		if (it != nullptr)
		{
			tlbMeta = {true, 1};
			return it;
		}
	}
	return nullptr;
	vpfn >>= 9;
	idx = vpfn & 0x0f, tag = vpfn >> 4;
	if (l1d_1g.find(idx) != l1d_1g.end())
	{
		auto it = l1d_1g[idx].get(tag);
		if (it != nullptr)
		{
			tlbMeta = {false, 2};
			return it;
		}
	}
	idx = vpfn & 0x7f, tag = vpfn >> 7;
	if (l2_1g.find(idx) != l2_1g.end())
	{
		auto it = l2_1g[idx].get(tag);
		if (it != nullptr)
		{
			tlbMeta = {true, 2};
			return it;
		}
	}
	return nullptr;
}

std::pair<bool, uint64_t> MMU::operator()(uint64_t *addr, uint8_t access, uint64_t stat, uint64_t *pbtr)
{
	if (stat ^ PageEn)
		return {true, 0};
	bool priv = (stat & Privilege);
	if (*addr < *(priv ? RvReg.first : RvReg.second))
		return {false, ERROR::PageFault};
	if (stat ^ TLB_En)
		return update(addr, access, stat, pbtr);
	uint64_t vpfn = *addr >> 12;
	uint8_t idx;
	if (access & X) // Instruct Access
	{
		if (access & R) // Read Only
		{
			auto it = InsSearch(vpfn);
			if (it != nullptr) // Found
			{
				PTE &p = (it->second).first;
				if ((p.meta & valid) == 0) // Invalid,read memory to check again
					return {false, ERROR::PageFault};
				if ((p.meta & access) != access)
					return {false, ERROR::AccessFault};
				switch (tlbMeta.second)
				{
				case 0: // 4kiB
					if (tlbMeta.first)
					{
						idx = vpfn & 0x0f;
						l1i_4k[idx].put((vpfn >> 4), {(vpfn >> 4), it->second});
					}
					*addr &= 0xfff;
					*addr += p.pfn << 12;
					return {true, 0};
				case 1: // 2MiB
					if (tlbMeta.first)
					{
						vpfn >>= 9;
						idx = vpfn & 0x0f;
						l1i_2m[idx].put((vpfn >> 4), {(vpfn >> 4), it->second});
					}
					*addr &= 0x1fffff;
					*addr += p.pfn << 12;
					return {true, 0};
				case 2: // 1GiB
					*addr &= 0x3fffffff;
					*addr += p.pfn << 12;
					return {true, 0};
				}
				return {true, 0};
			}
			else
				update(addr, access, stat, pbtr);
		}
		else // Write,Access Fault
			return {false, ERROR::AccessFault};
	}
	else // Data Access
	{
		auto it = DataSearch(vpfn);
		if (it != nullptr)
		{
			PTE &p = (it->second).first;
			if ((p.meta & valid) == 0)
				return {false, ERROR::PageFault};
			if ((p.meta & access) != access)
				return {false, ERROR::AccessFault};
			switch (tlbMeta.second)
			{
			case 0: // 4kiB
				if (tlbMeta.first)
				{
					idx = vpfn & 0x0f;
					l1d_4k[idx].put((vpfn >> 4), {(vpfn >> 4), it->second});
				}
				if (access & W)
					p.meta |= uint16_t(PMETA::DIRTY);
				*addr &= 0xfff;
				*addr += p.pfn << 12;
				return {true, 0};
			case 1: // 2MiB
				if (tlbMeta.first)
				{
					vpfn >>= 9;
					idx = vpfn & 0x07;
					l1d_2m[idx].put((vpfn >> 3), {(vpfn >> 3), it->second});
				}
				if (access & W)
					p.meta |= uint16_t(PMETA::DIRTY);
				*addr &= 0x1fffff;
				*addr += p.pfn << 12;
				return {true, 0};
			case 2: // 1GiB
				if (tlbMeta.first)
				{
					vpfn >>= 18;
					idx = vpfn & 0x07;
					l1d_1g[idx].put((vpfn >> 3), {(vpfn >> 3), it->second});
				}
				if (access & W)
					p.meta |= uint16_t(PMETA::DIRTY);
				*addr &= 0x3fffffff;
				*addr += p.pfn << 12;
				return {true, 0};
			}
		}
		else
			return update(addr, access, stat, pbtr);
	}
	return {true, 0};
}
std::pair<bool, uint64_t> MMU::update(uint64_t *addr, uint8_t access, uint64_t stat, uint64_t *pbtr)
{
	uint64_t vpfn = *addr >> 12;
	std::pair<uint32_t, uint16_t> splits[] = {
		{vpfn >> 36, (vpfn >> 27) & 0x1ff},
		{vpfn >> 27, (vpfn >> 18) & 0x1ff},
		{vpfn >> 18, (vpfn >> 9) & 0x1ff},
		{vpfn >> 9, vpfn & 0x1ff}};
	int i = 0;
	bool start = true;
	for (i = 0; i < 4; i++)
		if (splits[i].second > 0)
			break;
	uint64_t _addr, _page;
	PTE temp;
	for (; i < 4; i++)
	{
		_addr = start ? pbtr[i] : _addr;
		_addr += uint64_t(splits[i].second) << 3;
		start = false;
		global_ram[_addr >> 6].r_lock();
		memcpy(&_page, global_ram[_addr >> 6].get().data() + (_addr & 0x3f), 8);
		global_ram[_addr >> 6].r_unlock();
		temp = PTE(_page);
		if (temp.meta & valid) // we have found a valid page
		{
			switch (i)
			{
			case 0: // 512GiB Page,Invalid
				return {false, ERROR::PageFault};
			case 1: // 1GiB Page
				if (access & X)
					// Instruction fetch should not occur 1GiB pages
					return {false, ERROR::PageFault};
				if (access & W)
					temp.meta |= uint16_t(PMETA::DIRTY);
				if (stat & TLB_En)
				{ // update a TLB item
					Item it = {temp, _addr};
					uint64_t key = splits[2].first;
					PageList p = {key >> 6, it};
					uint8_t idx = key & 0x3f;
					l2_1g[idx].put(key >> 6, p);
				}
				*addr &= 0x3fffffff;
				*addr += temp.pfn << 12;
				return {true, 0};
			case 2: // 2MiB Page
				if (access & W)
					temp.meta |= uint16_t(PMETA::DIRTY);
				if (stat & TLB_En)
				{ // update a TLB item
					Item it = {temp, _addr};
					uint64_t key = splits[3].first;
					PageList p = {key >> 6, it};
					uint8_t idx = key & 0x3f;
					l2_2m[idx].put(key >> 6, p);
				}
				*addr &= 0x1fffff;
				*addr += temp.pfn << 12;
				return {true, 0};
			case 3: // 4KiB Page
				if (access & W)
					temp.meta |= uint16_t(PMETA::DIRTY);
				if (stat & TLB_En)
				{ // update a TLB item
					Item it = {temp, _addr};
					uint64_t key = vpfn;
					PageList p = {key >> 7, it};
					uint8_t idx = key & 0x7f;
					l2_4k[idx].put(key >> 7, p);
				}
				*addr &= 0xfff;
				*addr += temp.pfn << 12;
				return {true, 0};
			}
		}
	}
	return {false, ERROR::PageFault};
}
void MMU::flush()
{
	for (auto &i : l1d_4k)
		i.second.clear();
	for (auto &i : l1d_2m)
		i.second.clear();
	for (auto &i : l1d_1g)
		i.second.clear();
	for (auto &i : l2_4k)
		i.second.clear();
	for (auto &i : l2_2m)
		i.second.clear();
	for (auto &i : l2_1g)
		i.second.clear();
}